from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 加载采样数据
df = pd.read_csv("data.csv", sep=',', dtype=float, header=None)
y = df.iloc[0,:].values
X = []

# 采样时间秒，作为x轴
dt = 0.006
for i in range(len(y)):
    X.append(dt * i)

X = np.array(X, dtype=float).reshape(-1, 1)
y = y.reshape(-1, 1)

# print(df)
print("X and y Data shape: ", X.shape, y.shape)


# 拟合曲线
model = LinearRegression()
model.fit(X, y)
y_pred = model.predict(X)

# 保存参数 
k = model.coef_[0][0]
b = model.intercept_[0]

print("r2: ", r2_score(y, y_pred))
print(f"y = {k} * x + {b}")

f = open("model.txt", "w")
f.write(f"k = {k:.6f}, b = {b:.6f}\n")
f.close()

plt.plot(X, y)
plt.title("Raw curve")
plt.show()